package graph;

import java.util.Comparator;
import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;

/**
 * @author Liaorun
 */
public class Prim {


    public static class EdgeComparator implements Comparator<Edge> {

        @Override
        public int compare(Edge o1, Edge o2) {
            return o1.weight - o2.weight;
        }
    }


    public static Set<Edge> primMinimumSpanningTree(Graph graph) {
        // 解锁的边进入小顶堆
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(new EdgeComparator());

        HashSet<Node> set = new HashSet<>();

        // 记录最后需要的边的无序集合
        Set<Edge> result = new HashSet<>();

        // 随便从那个点开始都可以
        for (Node node : graph.nodes.values()) {
            // node 是开始点
            if (!set.contains(node)) {
                set.add(node);

                // 解锁所有相连的边
                priorityQueue.addAll(node.edges);

                while (!priorityQueue.isEmpty()) {
                    // 获得解锁的边中权值最小的一个
                    Edge edge = priorityQueue.poll();
                    // 可能的一个新的点
                    Node toNode = edge.to;
                    // 集合不含有的就是新的点
                    if (!set.contains(toNode)) {
                        // 记录新的点
                        set.add(toNode);
                        // 记录需要的边
                        result.add(edge);

                        // 解锁新的点所有相连的边
                        // 对于无向图，边回重复进入堆，但是没关系，这些边连接的点都已经在set集合中
                        // 不会影响算法的结果
                        priorityQueue.addAll(toNode.edges);
                    }
                }
            }

            // 如果是连通图就要break; 如果是森林就不能要
            break;
        }

        return result;
    }
}
